#!/bin/bash
# Linux 环境快速安装脚本（简化版本）

set -e  # 遇到错误立即退出

echo "========================================="
echo "  多无人机轨迹规划 - Linux 环境安装"
echo "========================================="
echo ""

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color

# 检查 conda 是否安装
if ! command -v conda &> /dev/null; then
    echo -e "${RED}错误: 未检测到 conda${NC}"
    echo "请先安装 Miniconda 或 Anaconda："
    echo "  wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh"
    echo "  bash Miniconda3-latest-Linux-x86_64.sh"
    exit 1
fi

echo -e "${GREEN}✓ 检测到 conda${NC}"

# 询问是否使用 GPU
read -p "是否安装 CUDA 版本的 PyTorch? (y/n): " -n 1 -r
echo
USE_GPU=$REPLY

# 创建 conda 环境
echo ""
echo "步骤 1/7: 创建基础 Python 环境..."
if conda env list | grep -q "^vae "; then
    echo -e "${YELLOW}环境 'vae' 已存在，是否删除并重新创建? (y/n)${NC}"
    read -p "" -n 1 -r
    echo
    if [[ $REPLY =~ ^[Yy]$ ]]; then
        conda env remove -n vae -y
    else
        echo "将使用现有环境..."
    fi
fi

# 使用最小化方案避免依赖冲突
echo "创建 Python 3.10.9 环境..."
conda create -n vae python=3.10.9 -y

echo -e "${GREEN}✓ 基础环境创建完成${NC}"

# 激活环境
echo ""
echo "步骤 2/7: 激活环境..."
source $(conda info --base)/etc/profile.d/conda.sh
conda activate vae
echo -e "${GREEN}✓ 环境已激活${NC}"

# 安装科学计算核心库（conda 安装性能更好）
echo ""
echo "步骤 3/7: 安装科学计算核心库..."
conda install -c conda-forge numpy scipy matplotlib -y
echo -e "${GREEN}✓ 科学计算库安装完成${NC}"

# 安装优化求解器（必须用 conda）
echo ""
echo "步骤 4/7: 安装优化求解器..."
conda install -c conda-forge casadi ipopt tqdm -y
echo -e "${GREEN}✓ 优化求解器安装完成${NC}"

# 安装 PyTorch
echo ""
echo "步骤 5/7: 安装 PyTorch..."
if [[ $USE_GPU =~ ^[Yy]$ ]]; then
    # 检测 CUDA 版本
    if command -v nvidia-smi &> /dev/null; then
        CUDA_VERSION=$(nvidia-smi | grep "CUDA Version" | awk '{print $9}' | cut -d. -f1,2)
        echo "检测到 CUDA 版本: $CUDA_VERSION"
        
        if [[ "$CUDA_VERSION" > "12.0" ]]; then
            echo "安装 CUDA 12.1 版本的 PyTorch..."
            pip install torch==2.8.0 torchvision==0.23.0 --index-url https://download.pytorch.org/whl/cu121
        elif [[ "$CUDA_VERSION" > "11.7" ]]; then
            echo "安装 CUDA 11.8 版本的 PyTorch..."
            pip install torch==2.8.0 torchvision==0.23.0 --index-url https://download.pytorch.org/whl/cu118
        else
            echo -e "${YELLOW}警告: CUDA 版本过低，安装 CPU 版本${NC}"
            pip install torch==2.8.0 torchvision==0.23.0 --index-url https://download.pytorch.org/whl/cpu
        fi
    else
        echo -e "${YELLOW}未检测到 NVIDIA GPU，安装 CPU 版本${NC}"
        pip install torch==2.8.0 torchvision==0.23.0 --index-url https://download.pytorch.org/whl/cpu
    fi
else
    echo "安装 CPU 版本的 PyTorch..."
    pip install torch==2.8.0 torchvision==0.23.0 --index-url https://download.pytorch.org/whl/cpu
fi
echo -e "${GREEN}✓ PyTorch 安装完成${NC}"

# 安装其他 Python 依赖
echo ""
echo "步骤 6/7: 安装其他依赖..."
pip install jax jaxlib julia wandb tensorboard pandas h5py seaborn pydantic pyyaml click requests pillow
pip install -e .
echo -e "${GREEN}✓ 依赖安装完成${NC}"

# 安装 Julia
echo ""
echo "步骤 7/7: 检查 Julia 安装..."
if command -v julia &> /dev/null; then
    JULIA_VERSION=$(julia --version | awk '{print $3}')
    echo -e "${GREEN}✓ 检测到 Julia $JULIA_VERSION${NC}"
else
    echo -e "${YELLOW}未检测到 Julia，需要手动安装:${NC}"
    echo "  1. 下载: wget https://julialang-s3.julialang.org/bin/linux/x64/1.9/julia-1.9.4-linux-x86_64.tar.gz"
    echo "  2. 解压: tar -xzf julia-1.9.4-linux-x86_64.tar.gz"
    echo "  3. 移动: sudo mv julia-1.9.4 /opt/"
    echo "  4. 链接: sudo ln -s /opt/julia-1.9.4/bin/julia /usr/local/bin/julia"
    echo "  5. 配置: echo 'export JULIA_BINDIR=/opt/julia-1.9.4/bin' >> ~/.bashrc"
    echo ""
    read -p "按 Enter 继续（假设你稍后会安装 Julia）..."
fi

# 配置 PyJulia
echo ""
echo "步骤 5/6: 配置 PyJulia..."
if command -v julia &> /dev/null; then
    python -c "import julia; julia.install()" || echo -e "${YELLOW}PyJulia 配置失败，请稍后手动配置${NC}"
    echo -e "${GREEN}✓ PyJulia 配置完成${NC}"
else
    echo -e "${YELLOW}跳过 PyJulia 配置（Julia 未安装）${NC}"
fi

# 运行验证
echo ""
echo "========================================="
echo "  运行环境验证..."
echo "========================================="
python verify_environment.py

# 完成
echo ""
echo "========================================="
echo -e "${GREEN}  安装完成！${NC}"
echo "========================================="
echo ""
echo "使用方法："
echo "  1. 激活环境: conda activate vae"
echo "  2. 运行验证: python verify_environment.py"
echo "  3. 生成数据: python examples/generate_data.py"
echo "  4. 训练模型: python examples/train_vae.py"
echo ""
echo "如果遇到问题，请查看 MIGRATION_GUIDE.md"
echo ""
